#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
妙趣拾光锦盒用户系统 - Gunicorn优化脚本
自动分析系统资源并生成最优的Gunicorn配置
"""

import os
import sys
import platform
import psutil
import argparse
import subprocess
import shutil
import socket
from pathlib import Path

# 颜色定义
GREEN = '\033[0;32m'
RED = '\033[0;31m'
YELLOW = '\033[0;33m'
NC = '\033[0m'  # 无颜色

def log_info(message):
    """输出信息日志"""
    print(f"{GREEN}[INFO]{NC} {message}")

def log_warn(message):
    """输出警告日志"""
    print(f"{YELLOW}[WARN]{NC} {message}")

def log_error(message):
    """输出错误日志"""
    print(f"{RED}[ERROR]{NC} {message}")

def check_root():
    """检查是否为root用户"""
    if os.geteuid() != 0:
        log_error("此脚本必须以root权限运行")
        sys.exit(1)

def detect_system_resources():
    """检测系统资源"""
    # CPU核心数
    cpu_count = psutil.cpu_count(logical=True)
    # 物理内存(GB)
    memory_gb = psutil.virtual_memory().total / (1024 * 1024 * 1024)
    # 磁盘空间
    disk_gb = psutil.disk_usage('/').total / (1024 * 1024 * 1024)
    
    # 获取操作系统信息
    os_name = platform.system()
    os_version = platform.release()
    
    # 获取网络接口
    network_interfaces = list(psutil.net_if_addrs().keys())
    
    # 获取主机名
    hostname = socket.gethostname()
    
    # 系统位数
    bits = platform.architecture()[0]
    
    # 获取系统负载
    try:
        load_avg = os.getloadavg()
    except:
        load_avg = (0, 0, 0)
    
    return {
        'cpu_count': cpu_count,
        'memory_gb': memory_gb,
        'disk_gb': disk_gb,
        'os_name': os_name,
        'os_version': os_version,
        'hostname': hostname,
        'bits': bits,
        'load_avg': load_avg,
        'network_interfaces': network_interfaces
    }

def calculate_gunicorn_settings(resources):
    """根据系统资源计算最优Gunicorn设置"""
    # 工作进程数计算 (2*CPU核心数+1)是一个常用的经验公式
    workers = 2 * resources['cpu_count'] + 1
    
    # 设置上限，避免过多workers
    if workers > 12:
        workers = 12
    
    # 根据内存计算每个worker的线程数
    threads = 2
    if resources['memory_gb'] > 8:
        threads = 4
    if resources['memory_gb'] > 16:
        threads = 8
    
    # 计算worker类型，根据内存选择合适的worker类
    worker_class = "gevent"
    
    # 计算超时设置
    timeout = 120
    
    # 计算最大请求数
    max_requests = 1000
    max_requests_jitter = 100
    
    # 计算keepalive值
    keepalive = 5
    
    return {
        'workers': workers,
        'threads': threads,
        'worker_class': worker_class,
        'timeout': timeout,
        'max_requests': max_requests,
        'max_requests_jitter': max_requests_jitter,
        'keepalive': keepalive
    }

def generate_gunicorn_config(config_path, settings):
    """生成Gunicorn配置文件"""
    config_content = f"""# Gunicorn配置文件 - 妙趣拾光锦盒用户系统
# 自动生成于 {os.path.basename(__file__)}
import multiprocessing
import os

# 绑定地址和端口 - 绑定到所有接口，允许从任何IP访问
bind = "0.0.0.0:8000"

# 工作进程数量
workers = {settings['workers']}  # 2*CPUs+1

# 工作线程数量
threads = {settings['threads']}

# 工作模式 - 使用gevent异步处理请求
worker_class = "{settings['worker_class']}"

# 请求超时时间(秒)
timeout = {settings['timeout']}

# 优雅重启超时(秒)
graceful_timeout = 30

# 最大请求数，超过此数量后worker会重启
max_requests = {settings['max_requests']}
max_requests_jitter = {settings['max_requests_jitter']}

# 数据包大小(字节)
limit_request_line = 4096
limit_request_fields = 100
limit_request_field_size = 8190

# 日志配置
errorlog = os.path.join("/var/log/box-user", "gunicorn-error.log")
accesslog = os.path.join("/var/log/box-user", "gunicorn-access.log")
loglevel = "info"
access_log_format = '%({X-Forwarded-For}i)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s "%(f)s" "%(a)s"'

# 守护进程模式 (在supervisor下不应设置为True)
daemon = False

# 进程名称
proc_name = "box-user"

# 保持连接选项
keepalive = {settings['keepalive']}

# 优雅地关闭worker
preload_app = True

# 请求处理前钩子
def on_starting(server):
    server.log.info("Starting Gunicorn server for box-user application")

# 配置服务器关闭事件钩子
def on_exit(server):
    server.log.info("Stopping Gunicorn server for box-user application")
"""
    # 创建目录（如果不存在）
    os.makedirs(os.path.dirname(config_path), exist_ok=True)
    
    # 写入配置文件
    with open(config_path, 'w') as f:
        f.write(config_content)
    
    return config_path

def update_start_script(script_path, app_dir, config_path):
    """更新应用启动脚本"""
    script_content = f"""#!/bin/bash
cd {app_dir}
source venv/bin/activate
export DJANGO_SETTINGS_MODULE=boxuser.settings_production
exec gunicorn --config={config_path} boxuser.wsgi:application
"""
    # 创建目录（如果不存在）
    os.makedirs(os.path.dirname(script_path), exist_ok=True)
    
    # 写入启动脚本
    with open(script_path, 'w') as f:
        f.write(script_content)
    
    # 添加执行权限
    os.chmod(script_path, 0o755)
    
    return script_path

def update_supervisor_config(supervisor_config, start_script, app_dir):
    """更新Supervisor配置"""
    config_content = f"""[program:box-user]
command={start_script}
directory={app_dir}
user=root
autostart=true
autorestart=true
startretries=3
stopasgroup=true
killasgroup=true
stderr_logfile=/var/log/box-user/supervisor.err.log
stdout_logfile=/var/log/box-user/supervisor.out.log
environment=LANG=en_US.UTF-8,LC_ALL=en_US.UTF-8
"""
    # 备份当前配置（如果存在）
    if os.path.exists(supervisor_config):
        shutil.copy2(supervisor_config, f"{supervisor_config}.bak")
    
    # 创建目录（如果不存在）
    os.makedirs(os.path.dirname(supervisor_config), exist_ok=True)
    
    # 写入配置
    with open(supervisor_config, 'w') as f:
        f.write(config_content)
    
    return supervisor_config

def install_dependencies(app_dir):
    """安装必要的Python依赖"""
    try:
        # 激活虚拟环境并安装依赖
        activate_cmd = f"cd {app_dir} && source venv/bin/activate && pip install gunicorn gevent"
        subprocess.run(activate_cmd, shell=True, check=True)
        log_info("依赖安装成功")
        return True
    except subprocess.CalledProcessError:
        log_error("依赖安装失败")
        return False

def restart_service():
    """重启应用服务"""
    try:
        # 重新加载Supervisor配置
        subprocess.run("supervisorctl update box-user", shell=True, check=True)
        # 重启应用
        subprocess.run("supervisorctl restart box-user", shell=True, check=True)
        log_info("应用服务已重启")
        return True
    except subprocess.CalledProcessError:
        log_error("应用服务重启失败")
        return False

def check_service_status():
    """检查服务状态"""
    try:
        result = subprocess.run("supervisorctl status box-user", shell=True, 
                                capture_output=True, text=True, check=True)
        status = result.stdout
        print(status)
        
        if "RUNNING" in status:
            log_info("应用服务运行正常")
            return True
        else:
            log_warn("应用服务可能未正常运行")
            return False
    except subprocess.CalledProcessError:
        log_error("无法获取服务状态")
        return False

def test_service_accessibility():
    """测试服务可访问性"""
    endpoints = [
        "http://localhost:8000/",
        "http://127.0.0.1:8000/"
    ]
    
    # 获取服务器IP
    try:
        server_ip = socket.gethostbyname(socket.gethostname())
        endpoints.append(f"http://{server_ip}:8000/")
    except:
        log_warn("无法获取服务器IP地址")
    
    success = True
    
    for endpoint in endpoints:
        try:
            cmd = f"curl -s -o /dev/null -w \"{endpoint} HTTP状态码: %{{http_code}}\\n\" {endpoint}"
            subprocess.run(cmd, shell=True, check=True)
        except subprocess.CalledProcessError:
            log_warn(f"无法访问 {endpoint}")
            success = False
    
    return success

def ensure_log_directory(log_dir):
    """确保日志目录存在"""
    try:
        os.makedirs(log_dir, exist_ok=True)
        os.chmod(log_dir, 0o755)
        log_info(f"已创建/确认日志目录: {log_dir}")
        return True
    except Exception as e:
        log_error(f"创建日志目录失败: {e}")
        return False

def update_nginx_config(nginx_config):
    """更新Nginx配置"""
    if not os.path.exists(nginx_config):
        log_warn(f"Nginx配置文件不存在: {nginx_config}")
        return False
    
    try:
        # 读取当前配置
        with open(nginx_config, 'r') as f:
            content = f.read()
        
        # 检查是否需要更新
        if "proxy_pass http://127.0.0.1:8000" in content:
            # 备份配置
            shutil.copy2(nginx_config, f"{nginx_config}.bak")
            
            # 更新配置
            content = content.replace("proxy_pass http://127.0.0.1:8000", "proxy_pass http://localhost:8000")
            
            # 写入更新后的配置
            with open(nginx_config, 'w') as f:
                f.write(content)
            
            # 重新加载Nginx配置
            subprocess.run("systemctl reload nginx", shell=True, check=True)
            
            log_info("Nginx配置已更新")
        else:
            log_info("Nginx配置无需更新")
        
        return True
    except Exception as e:
        log_error(f"更新Nginx配置失败: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='妙趣拾光锦盒用户系统 - Gunicorn优化脚本')
    parser.add_argument('--app-dir', default='/var/www/box-user', help='应用目录路径')
    parser.add_argument('--supervisor-config', default='/etc/supervisord.d/box-user.ini', help='Supervisor配置文件路径')
    parser.add_argument('--nginx-config', default='/etc/nginx/conf.d/box-user.conf', help='Nginx配置文件路径')
    parser.add_argument('--log-dir', default='/var/log/box-user', help='日志目录路径')
    parser.add_argument('--dry-run', action='store_true', help='只显示将要进行的操作，不实际执行')
    
    args = parser.parse_args()
    
    # 检查root权限
    check_root()
    
    # 显示脚本信息
    log_info("妙趣拾光锦盒用户系统 - Gunicorn优化脚本")
    log_info("=" * 50)
    
    # 检测系统资源
    log_info("正在分析系统资源...")
    resources = detect_system_resources()
    
    log_info(f"系统信息:")
    log_info(f"  操作系统: {resources['os_name']} {resources['os_version']} ({resources['bits']})")
    log_info(f"  主机名: {resources['hostname']}")
    log_info(f"  CPU核心数: {resources['cpu_count']}")
    log_info(f"  内存大小: {resources['memory_gb']:.2f} GB")
    log_info(f"  磁盘空间: {resources['disk_gb']:.2f} GB")
    log_info(f"  系统负载: {resources['load_avg']}")
    
    # 计算Gunicorn设置
    log_info("计算最优Gunicorn设置...")
    settings = calculate_gunicorn_settings(resources)
    
    log_info(f"Gunicorn配置:")
    log_info(f"  Workers: {settings['workers']}")
    log_info(f"  Threads: {settings['threads']}")
    log_info(f"  Worker类: {settings['worker_class']}")
    log_info(f"  超时设置: {settings['timeout']}秒")
    log_info(f"  最大请求数: {settings['max_requests']}")
    
    if args.dry_run:
        log_info("Dry run模式 - 不执行实际操作")
        return 0
    
    # 检查应用目录
    if not os.path.exists(args.app_dir):
        log_error(f"应用目录不存在: {args.app_dir}")
        return 1
    
    # 确保日志目录存在
    ensure_log_directory(args.log_dir)
    
    # 生成Gunicorn配置文件
    gunicorn_config = os.path.join(args.app_dir, "gunicorn_config.py")
    log_info(f"生成Gunicorn配置文件: {gunicorn_config}")
    generate_gunicorn_config(gunicorn_config, settings)
    
    # 更新启动脚本
    start_script = os.path.join(args.app_dir, "start.sh")
    log_info(f"更新应用启动脚本: {start_script}")
    update_start_script(start_script, args.app_dir, gunicorn_config)
    
    # 安装必要依赖
    log_info("安装必要的Python依赖...")
    install_dependencies(args.app_dir)
    
    # 更新Supervisor配置
    log_info(f"更新Supervisor配置: {args.supervisor_config}")
    update_supervisor_config(args.supervisor_config, start_script, args.app_dir)
    
    # 更新Nginx配置
    log_info("检查Nginx配置...")
    update_nginx_config(args.nginx_config)
    
    # 重启服务
    log_info("重启应用服务...")
    restart_service()
    
    # 等待服务启动
    log_info("等待服务启动...")
    import time
    time.sleep(5)
    
    # 检查服务状态
    log_info("检查服务状态...")
    service_ok = check_service_status()
    
    if service_ok:
        # 测试服务可访问性
        log_info("测试服务可访问性...")
        test_service_accessibility()
        
        log_info("Gunicorn优化配置完成!")
        log_info(f"配置文件: {gunicorn_config}")
        log_info(f"启动脚本: {start_script}")
        log_info(f"Supervisor配置: {args.supervisor_config}")
        log_info("注意: Gunicorn现在绑定到0.0.0.0:8000，可以从任何IP地址访问")
    else:
        log_error("应用可能未正常启动，请检查日志:")
        subprocess.run(f"tail -n 50 {args.log_dir}/supervisor.err.log", shell=True)
    
    return 0

if __name__ == "__main__":
    sys.exit(main()) 